# Copyright 2025 Amazon Inc

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import time
from datetime import datetime, timezone

from playwright.sync_api import Page
from typing_extensions import TypedDict

from nova_act.tools.browser.default.util.image_helpers import (
    compare_images,
    take_screenshot_as_data_url,
)
from nova_act.tools.browser.default.util.take_observation import save_data_url_to_file
from nova_act.util.logging import setup_logging

_LOGGER = setup_logging(__name__)


def timed_wait(milliseconds: int, page: Page) -> None:
    page.wait_for_timeout(milliseconds)


class ConsecutiveChecksOptions(TypedDict):
    max_timeout_ms: int | None  # in milliseconds
    number_of_checks: int
    percent_difference_threshold: float
    polling_interval_ms: int  # in milliseconds


WAIT_FOR_PAGE_TO_SETTLE_CONFIG: ConsecutiveChecksOptions = {
    "max_timeout_ms": 3000,
    "number_of_checks": 3,
    "percent_difference_threshold": 25,
    "polling_interval_ms": 500,
}


class PreviousAttempt(TypedDict):
    attempt_number: int
    screenshot: str


def delay(milliseconds: int) -> None:
    """Helper function to delay execution for a specified number of milliseconds."""
    time.sleep(milliseconds / 1000)


def consecutive_identical_checks(
    page: Page,
    options: ConsecutiveChecksOptions,
    save_screenshots: bool = False,
) -> None:
    """
    This function checks if the page has been unchanged for a certain number of checks.
    It will return when either:
    1. The max timeout has been reached
    2. The page has been stable for the specified number of checks
    """

    start_time = datetime.now(timezone.utc)
    consecutive_stable_count = 0
    previous_screenshot = None

    while True:
        timestamp = datetime.now().isoformat()

        # Check timeout
        if (
            options["max_timeout_ms"] is not None
            and options["max_timeout_ms"] > 0
            and (datetime.now().timestamp() - start_time.timestamp()) * 1000 >= options["max_timeout_ms"]
        ):
            _LOGGER.debug(f"[CONSECUTIVE_CHECK] {timestamp} - TIMEOUT REACHED")
            return

        _LOGGER.debug(
            f"[CONSECUTIVE_CHECK] {timestamp} - Attempt {consecutive_stable_count}, "
            f"Threshold: {options['percent_difference_threshold']}%"
        )

        # Take screenshot
        screenshot = take_screenshot_as_data_url(page)
        if screenshot is None:
            _LOGGER.debug(f"[CONSECUTIVE_CHECK] {timestamp} - SCREENSHOT FAILED")
            raise RuntimeError("Attempted to take a screenshot, but failed. Please try again.")

        if save_screenshots:
            save_data_url_to_file(screenshot, f"screenshot_{consecutive_stable_count}.jpeg")

        # First iteration - just wait and continue
        if previous_screenshot is None:
            _LOGGER.debug(
                f"[CONSECUTIVE_CHECK] {timestamp} - FIRST ATTEMPT, waiting {options['polling_interval_ms']}ms"
            )
            previous_screenshot = screenshot
            delay(options["polling_interval_ms"])
            continue

        # Compare with previous screenshot
        percent_difference = compare_images(previous_screenshot, screenshot)
        _LOGGER.debug(
            f"[CONSECUTIVE_CHECK] {timestamp} - Attempt {consecutive_stable_count}: "
            f"{percent_difference:.2f}% difference"
        )

        # If difference is too high, reset counter
        if percent_difference >= options["percent_difference_threshold"]:
            _LOGGER.debug(f"[CONSECUTIVE_CHECK] {timestamp} - DIFFERENCE TOO HIGH, resetting to attempt 1")
            consecutive_stable_count = 0
        else:
            consecutive_stable_count += 1

            # Check if we've reached required stable checks
            if consecutive_stable_count >= options["number_of_checks"]:
                _LOGGER.debug(
                    f"[CONSECUTIVE_CHECK] {timestamp} - COMPLETED after {consecutive_stable_count} stable checks"
                )
                return

        previous_screenshot = screenshot
        delay(options["polling_interval_ms"])


def wait_for_page_to_settle(
    page: Page,
    options: ConsecutiveChecksOptions,
    save_screenshots: bool = False,
) -> None:
    try:
        page.wait_for_load_state("load")
    except Exception as e:
        _LOGGER.debug(f"Failed to wait for page to load. Error: {e}")
        pass
    consecutive_identical_checks(page, options, save_screenshots)
